#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>

// 全局变量：无人机的当前状态
mavros_msgs::State current_state;

// 状态回调函数：更新当前无人机状态
void state_cb(const mavros_msgs::State::ConstPtr& msg) {
    current_state = *msg;
}

int main(int argc, char **argv) {
    ros::init(argc, argv, "offboard_node");
    ros::NodeHandle nh;

    // 订阅无人机的状态
    ros::Subscriber state_sub = nh.subscribe<mavros_msgs::State>("/mavros/state", 10, state_cb);

    // 发布目标位置的发布者
    ros::Publisher local_pos_pub = nh.advertise<geometry_msgs::PoseStamped>("/mavros/setpoint_position/local", 10);

    // 服务客户端：用于解锁无人机和设置模式
    ros::ServiceClient arming_client = nh.serviceClient<mavros_msgs::CommandBool>("/mavros/cmd/arming");
    ros::ServiceClient set_mode_client = nh.serviceClient<mavros_msgs::SetMode>("/mavros/set_mode");

    // 设定控制频率为20Hz
    ros::Rate rate(20);

    // 等待与飞控的连接
    while (ros::ok() && !current_state.connected) {
        ros::spinOnce();
        rate.sleep();
    }

    ROS_INFO("Connected to PX4.");

    // 初始位置设置为(0, 0, 2)，即起飞到2米高度
    geometry_msgs::PoseStamped pose;
    pose.pose.position.x = 0;
    pose.pose.position.y = 0;
    pose.pose.position.z = 2;

    // 发布初始的姿态指令，防止PX4因为没有接收到指令而无法切换到OFFBOARD模式
    for (int i = 100; ros::ok() && i > 0; --i) {
        local_pos_pub.publish(pose);
        ros::spinOnce();
        rate.sleep();
    }

    // 进入等待循环，直到遥控器切换到OFFBOARD模式并解锁无人机
    ROS_INFO("Waiting for OFFBOARD mode and arming...");
    while (ros::ok()) {
        // 检查当前模式是否为 OFFBOARD，并且是否已经解锁
        if (current_state.mode == "OFFBOARD" && current_state.armed) {
            ROS_INFO("OFFBOARD mode and arming confirmed. Proceeding with the mission...");
            break;  // 退出等待循环，开始执行后续任务
        }

        // 如果没有进入OFFBOARD模式，提示等待
        if (current_state.mode != "OFFBOARD") {
            ROS_INFO_THROTTLE(5, "Waiting for OFFBOARD mode from remote control...");
        }

        // 如果没有解锁，提示等待
        if (!current_state.armed) {
            ROS_INFO_THROTTLE(5, "Waiting for arming from remote control...");
        }

        ros::spinOnce();
        rate.sleep();
    }

    ros::Time last_request = ros::Time::now();

    // 控制逻辑：起飞到2米高，前飞2米，返回起点并降落
    while (ros::ok()) {
        // 1. 起飞到2米高
        pose.pose.position.x = 0;
        pose.pose.position.y = 0;
        pose.pose.position.z = 2;
        local_pos_pub.publish(pose);
        ROS_INFO("Ascending to 2 meters");
        ros::Duration(5.0).sleep();  // 等待起飞完成

        // 2. 往前飞2米
        pose.pose.position.x = 2;
        pose.pose.position.y = 0;
        pose.pose.position.z = 2;
        local_pos_pub.publish(pose);
        ROS_INFO("Moving forward 2 meters");
        ros::Duration(5.0).sleep();  // 等待无人机前飞完成

        // 2. 返回起点
        pose.pose.position.x = 0;
        pose.pose.position.y = 0;
        pose.pose.position.z = 2;
        local_pos_pub.publish(pose);
        ROS_INFO("Returning to starting point");
        ros::Duration(5.0).sleep();  // 等待无人机返回

        // 4. 降落
        pose.pose.position.x = 0;
        pose.pose.position.y = 0;
        pose.pose.position.z = 0;
        local_pos_pub.publish(pose);
        ROS_INFO("Landing");
        ros::Duration(5.0).sleep();  // 等待降落完成

        ros::spinOnce();
        rate.sleep();
    }

    return 0;
}
